# -*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

from pymongo import UpdateOne, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes, get_trading_dates
from util.database import DB_CONN
import pandas as pd
import numpy as np

"""
实现市盈率因子的计算和保存
"""


class SP_TTM_Factor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='pe')

    def dualdate(self, a):

        # print (a)
        if len(a.strip()) < 10:
            month = '0' + a.split('-')[1]
            pata = a.split('-')
            pata[1] = month
            final = '-'.join(pata)
            print (final)
            return (final)
        else:
            print (a)
            return a

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值，并保存到数据库中
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        dm = DataModule()
        dates = get_trading_dates(begin_date, end_date)
        codes = ['601808']
        update_requests = []
        for code in codes:
            data = DB_CONN['mainreport'].find({'code': code})
            cc = [zz for zz in data]
            s = [float(bb['gsjlr'][:-1]) for bb in cc]
            s.reverse()
            cc = [asa['report_date'].replace('/', '-')[:10] for asa in cc]
            cc = [self.dualdate(x) for x in cc]
            jrl=[DB_CONN['mainreport'].find({'code': code,'date':self.dualdate(x)})[0]['mgjzc'] for x in cc]

            total=[DB_CONN['basic'].find({'code': code,'date':self.dualdate(x)})[0]['total'] for x in cc]

            print (total)
            print (jrl)
            total.reverse()
            jrl.reverse()
            cc.reverse()
            ROE=[s[asa]/float(jrl[asa])*float(total[asa]) for asa in range(0,len(s))]

            print (ROE)
            growth = [(s[asa + 1] - ROE[asa]) / abs(ROE[asa]) for asa in range(0, len(ROE) - 1)]

            growth.insert(0, None)
            maplist = dict(zip(cc, growth))
            # cc.reverse()
            # s.reverse()
            for date in dates:
                for asa in cc:
                    # print (date,asa)
                    if date < asa:
                        continue
                    else:
                        reportdat = asa
                    break

                ROE = maplist[reportdat]
                basic = DB_CONN['basic'].find_one({'code': code, 'date': date})
                close = dm.get_k_data(code=code, begin_date=date, end_date=date)

                TTM = int(basic['totals'] * close['close'])
                ROE_TTM = (ROE  / TTM)

                update_requests.append(
                    UpdateOne(
                        {'code': code, 'date': date},
                        {'$set': {'code': code, 'date': date, 'ROE_TTM': ROE_TTM}}, upsert=True))
                print(update_requests)
            #
            #     break
            # break


SP_TTM_Factor().compute('2018-01-02', '2018-09-05')
